// Copyright (c) 2022 Presto Labs Pte. Ltd.
// Author: chunhui

#include "coin2/strategy/util/top_percentile.h"

#include <algorithm>
#include <deque>
#include <random>
#include <unordered_map>
#include <utility>
#include <vector>

#include <gtest/gtest.h>

namespace coin2::strategy::util {

class SlowTopPercentile {
 public:
  explicit SlowTopPercentile(double percentile) : percentile_(percentile) {
    CHECK_GT(percentile, 0);
  }

  void Add(int raw_index, double value) {
    if (index_map_.find(raw_index) == index_map_.end()) {
      int id = static_cast<int>(val_.size());
      index_map_[raw_index] = id;
      val_.emplace_back(value);
    } else {
      val_[index_map_[raw_index]] = value;
    }
  }

  bool Get(double* value) {
    if (val_.size() == 1) {
      *value = val_[0];
      return true;
    }
    std::vector<double> tmp(val_.begin(), val_.end());
    std::sort(tmp.begin(), tmp.end());
    int size = tmp.size() - 1;
    auto id = static_cast<size_t>(size * percentile_);
    double x1 = static_cast<double>(id) / size;
    double x2 = static_cast<double>(id + 1) / size;
    double y1 = tmp[id];
    double y2 = tmp[id + 1];
    *value = y1 + (percentile_ - x1) * (y2 - y1) / (x2 - x1);
    return true;
  }

  void Debug() {
    std::stringstream ss;
    std::vector<double> tmp(val_.begin(), val_.end());
    std::sort(tmp.begin(), tmp.end());
    for (int i = 0; i < tmp.size(); ++i) {
      ss << tmp[i] << ", ";
    }
    LOG(ERROR) << ss.str();
  }

 private:
  double percentile_;
  std::unordered_map<int, int> index_map_;
  std::vector<double> val_;
};

void GenerateTestCase(
    size_t test_size,
    int num_range,
    int index_range,
    std::vector<double>* value,
    std::vector<int>* index) {
  std::random_device device;
  std::mt19937 generator(device());
  std::uniform_int_distribution<> val_distribution(1, num_range);
  std::uniform_int_distribution<> idx_distribution(1, index_range);
  for (size_t i = 0; i < test_size; ++i) {
    int idx = idx_distribution(generator);
    index->push_back(idx);
    value->push_back(val_distribution(generator));
  }
}

TEST(TopPercentileTest, Basic) {
  size_t test_size = 15;
  double percentile = 0.5;
  TopPercentile tm(percentile);
  SlowTopPercentile sltm(percentile);
  std::vector<double>
      values{403, 545, 912, 187, 180, 390, 428, 212, 381, 737, 462, 807, 501, 148, 628};
  std::vector<int> indexes{54, 82, 73, 37, 46, 26, 98, 2, 94, 56, 40, 26, 51, 43, 43};
  for (int i = 0; i < test_size; ++i) {
    tm.Add(indexes[i], values[i]);
    sltm.Add(indexes[i], values[i]);
    double actual = 0;
    double expected = 1;
    tm.Get(&actual);
    sltm.Get(&expected);
    EXPECT_EQ(actual, expected);
  }
}

TEST(TopPercentileTest, Randomized) {
  int num_range = 10000;
  int index_range = 1000;
  size_t test_size = 100000;
  double percentile = 0.5;
  TopPercentile tm(percentile);
  SlowTopPercentile sltm(percentile);
  std::vector<double> values;
  std::vector<int> indexes;
  GenerateTestCase(test_size, num_range, index_range, &values, &indexes);

  for (int i = 0; i < test_size; ++i) {
    tm.Add(indexes[i], values[i]);
    sltm.Add(indexes[i], values[i]);
    double actual = 0;
    double expected = 1;
    tm.Get(&actual);
    sltm.Get(&expected);
    EXPECT_EQ(actual, expected);
  }
}

}  // namespace coin2::strategy::util
